#region License and Terms
// MoreLINQ - Extensions to LINQ to Objects
// Copyright (c) 2008 Jonathan Skeet. All rights reserved.
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//     http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#endregion

namespace MoreLinq
{
    using System;
    using System.Collections;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Linq;

    static partial class MoreEnumerable
    {
        private class BatchPartEnumerator<T> : IEnumerator<T>, IEnumerable<T>
        {
            private readonly IEnumerator<T> _etor;
            private bool _onFirst = true;
            private T _current;
            private int _count;
            public BatchPartEnumerator(IEnumerator<T> etor, T current, int count)
            {
                _etor = etor;
                _current = current;
                _count = count;
            }
            public T Current
            {
                get { return _current; }
            }
            object IEnumerator.Current
            {
                get { return _current; }
            }
            void IDisposable.Dispose()
            {
            }
            public bool MoveNext()
            {
                if(_onFirst)
                {
                    _onFirst = false;
                    return true;
                }
                if(_count == 0)
                    return false;
                if(_etor.MoveNext())
                {
                    _current = _etor.Current;
                    --_count;
                    return true;
                }
                return false;
            }
            public void EatUp()
            {
                while(_count != 0 && _etor.MoveNext())
                    --_count;
            }
            void IEnumerator.Reset()
            {
                throw new NotSupportedException();
            }
            public IEnumerator<T> GetEnumerator()
            {
                return this;
            }
            IEnumerator IEnumerable.GetEnumerator()
            {
                return this;
            }
        }
        private class BatchEnumerator<T> : IEnumerator<IEnumerable<T>>
        {
            private readonly IEnumerator<T> _etor;
            private readonly int _afterFirst;
            private T _current;
            private BatchPartEnumerator<T> _currentBatch;
            public BatchEnumerator(IEnumerator<T> etor, int size)
            {
                _etor = etor;
                _afterFirst = size - 1;
            }
            public IEnumerable<T> Current
            {
                get { return _currentBatch = new BatchPartEnumerator<T>(_etor, _current, _afterFirst); }
            }
            object IEnumerator.Current
            {
                get { return Current; }
            }
            public void Dispose()
            {
                _etor.Dispose();
            }
            public bool MoveNext()
            {
                if(_currentBatch != null)
                    _currentBatch.EatUp();
                if(!_etor.MoveNext())
                    return false;
                _current = _etor.Current;
                return true;
            }
            void IEnumerator.Reset()
            {
                throw new NotSupportedException();
            }
        }
        private class BatchEnumerable<T> : IEnumerable<IEnumerable<T>>
        {
            private readonly IEnumerable<T> _etor;
            private readonly int _size;
            public BatchEnumerable(IEnumerable<T> etor, int size)
            {
                _etor = etor;
                _size = size;
            }
            public IEnumerator<IEnumerable<T>> GetEnumerator()
            {
                return new BatchEnumerator<T>(_etor.GetEnumerator(), _size);
            }
            IEnumerator IEnumerable.GetEnumerator()
            {
                return GetEnumerator();
            }
        }
        /// <summary>
        /// Batches the source sequence into sized buckets.
        /// </summary>
        /// <typeparam name="TSource">Type of elements in <paramref name="source"/> sequence.</typeparam>
        /// <param name="source">The source sequence.</param>
        /// <param name="size">Size of buckets.</param>
        /// <returns>A sequence of equally sized buckets containing elements of the source collection.</returns>
        /// <remarks>
        /// This operator uses deferred execution and streams its results (buckets and bucket content). 
        /// It is also identical to <see cref="Partition{TSource}(System.Collections.Generic.IEnumerable{TSource},int)"/>.
        /// </remarks>
        public static IEnumerable<IEnumerable<TSource>> Batch<TSource>(this IEnumerable<TSource> source, int size)
        {
            if (source == null) throw new ArgumentNullException("source");
            if (size <= 0) throw new ArgumentOutOfRangeException("size");
            return new BatchEnumerable<TSource>(source, size);
        }

        /// <summary>
        /// Batches the source sequence into sized buckets and applies a projection to each bucket.
        /// </summary>
        /// <typeparam name="TSource">Type of elements in <paramref name="source"/> sequence.</typeparam>
        /// <typeparam name="TResult">Type of result returned by <paramref name="resultSelector"/>.</typeparam>
        /// <param name="source">The source sequence.</param>
        /// <param name="size">Size of buckets.</param>
        /// <param name="resultSelector">The projection to apply to each bucket.</param>
        /// <returns>A sequence of projections on equally sized buckets containing elements of the source collection.</returns>
        /// <remarks>
        /// This operator uses deferred execution and streams its results (buckets and bucket content).
        /// It is also identical to <see cref="Partition{TSource}(System.Collections.Generic.IEnumerable{TSource},int)"/>.
        /// </remarks>
        
        public static IEnumerable<TResult> Batch<TSource, TResult>(this IEnumerable<TSource> source, int size,
            Func<IEnumerable<TSource>, TResult> resultSelector)
        {
            if (source == null) throw new ArgumentNullException("source");
            if (size <= 0) throw new ArgumentOutOfRangeException("size");
            if (resultSelector == null) throw new ArgumentNullException("resultSelector");
            return new BatchEnumerable<TSource>(source, size).Select(resultSelector);
        }
    }
}
